#include "CutUtils.hxx"
#include "FiducialVolumeDefinition.hxx"
#include <TVector3.h>
#include <stdio.h>
#include <iostream>
#include <math.h>

//**************************************************
float cutUtils::GetSeparationSquared(const Float_t* pos1, const Float_t* pos2){
//**************************************************
    return ((pos1[0] - pos2[0])*(pos1[0] - pos2[0]) + (pos1[1] - pos2[1])*(pos1[1] - pos2[1]) + (pos1[2] - pos2[2])*(pos1[2] - pos2[2]));
}

//**************************************************
void cutUtils::AssignArray(Float_t* src, Float_t* dest, int n){
//**************************************************
    for(int i = 0; i < n; ++i){
        dest[i] = src[i];
    }
}

//**************************************************
void cutUtils::AssignArray(const Float_t* src, Float_t* dest, int n){
//**************************************************
    for(int i = 0; i < n; ++i){
        dest[i] = src[i];
    }
}

//**************************************************
bool cutUtils::TrackQualityCut(AnaTrackB& track){
//**************************************************

    if(track.TPCQualityCut > -1){
        return track.TPCQualityCut;
    }
    // Gets all segments in the closest TPC
    AnaSubTrackB* TPCSegment = anaUtils::GetSegmentWithMostNodesInClosestTpc(track);

    if (TPCSegment){
        bool passed = TPCTrackQualityCut(*TPCSegment);
        track.TPCQualityCut = passed;
        return passed;
    }
    else{
        track.TPCQualityCut = 0;
        return false;
    }
}

//**************************************************
bool cutUtils::TPCTrackQualityCut(const AnaSubTrackB& tpcSeg){
//**************************************************

    Float_t cut = 18;  
    if (tpcSeg.NNodes > cut) 
        return true;
    else
        return false;
}

//**************************************************
bool cutUtils::MuonPIDCut(const AnaTrackB& track, bool prod5Cut){
//**************************************************

    Float_t cut1 = 0.8;
    Float_t cut2 = 0.05;
    Float_t pcut = 500;

    Float_t PIDLikelihood[4];
    anaUtils::GetPIDLikelihood(track, PIDLikelihood, prod5Cut);

    if (((PIDLikelihood[0]+PIDLikelihood[3])/(1-PIDLikelihood[2]) > cut1 || track.Momentum > pcut ) && (PIDLikelihood[0]>cut2)){
        return true; 
    }

    return false;
}

//**************************************************
bool cutUtils::AntiMuonPIDCut(const AnaTrackB& track){
//**************************************************

    Float_t cut1 = 0.9;
    Float_t cut2 = 0.1;
    Float_t cut3 = 0.7;
    Float_t pcut = 500;

    Float_t PIDLikelihood[4];
    anaUtils::GetPIDLikelihood(track, PIDLikelihood);

    if (((PIDLikelihood[0]+PIDLikelihood[3])/(1-PIDLikelihood[2]) > cut1 || track.Momentum > pcut ) && (PIDLikelihood[0]>cut2 && PIDLikelihood[0]<cut3)){
        return true; 
    }

    return false;
}


//**************************************************
bool cutUtils::ProtonPIDCut(const AnaTrackB& track){
//**************************************************

    if (anaUtils::GetPIDLikelihood(track,2)>0.9)
        return true;
    else    
        return false;
}

//**************************************************
bool cutUtils::FiducialCut(const AnaTrackB& track, const SubDetId::SubDetEnum det){
//**************************************************
    if (!SubDetId::GetDetectorUsed(track.Detector, det)) return false;
    return FiducialCut(track.PositionStart, det);
}

//**************************************************
bool cutUtils::FiducialCut(const Float_t* pos, const SubDetId::SubDetEnum det){
//**************************************************

    return anaUtils::InFiducialVolume(det,pos);
}

//********************************************************************
bool cutUtils::DeltaZCut(const AnaTrackB& track){
//********************************************************************

    Float_t startposz = track.PositionStart[2];
    Float_t endposz   = track.PositionEnd[2];
    TVector3 vect(track.PositionStart[0], track.PositionStart[1], track.PositionStart[2]);
    if (anaUtils::InFiducialVolume(SubDetId::kFGD1, track.PositionStart) && endposz < startposz){
        return false;
    }
    else{
        return true;
    }
}

//********************************************************************
AnaTrackB* cutUtils::FindVetoTrack(const AnaEventB& event, const AnaTrackB& track){
//********************************************************************

    Float_t Hmom = -9999; 
    AnaTrackB* vetoTrack = NULL;
    AnaTrackB* selTracks[NMAXTRACKS];
    int nTPC = anaUtils::GetAllTracksUsingTPC(event, selTracks);

    //loop over tpc tracks
    for (Int_t i=0;i<nTPC; ++i){
        AnaTrackB* track_temp = selTracks[i];
        Float_t mom = track_temp->Momentum;

        if (&track == track_temp) continue;  // same particle as muon candidate 
        if( mom < 0 ) continue;   // protection for bad momentum
        if( mom > Hmom ) {        // Search for highest momentum track in the bunch that is not he muon. 
            vetoTrack =  track_temp;
            Hmom = mom;
        }
    }

    return vetoTrack;
}

//********************************************************************
bool cutUtils::ExternalVetoCut(const AnaTrackB& candidate, AnaTrackB* vetoTrack){
//********************************************************************

    Float_t TPCVetoDistance = -150.;

    Float_t startposz = candidate.PositionStart[2];
    Float_t endposz = candidate.PositionEnd[2];
    TVector3 vect(candidate.PositionStart[0], candidate.PositionStart[1], candidate.PositionStart[2]);
    // cut to handle a bug in prod 4 (115 is the limit of FGD1), should be ininfluent in prod5
    if(endposz < 115. && endposz < startposz && anaUtils::InFiducialVolume(SubDetId::kFGD2, candidate.PositionStart) ) return false;

    if(vetoTrack) {

        if(vetoTrack->Original){
            if(vetoTrack->Original->PositionStart[2] - candidate.PositionStart[2] < TPCVetoDistance)  return false; 
        }
    }

    return true;
}

//********************************************************************
bool cutUtils::ExternalVetoFGD2Cut(const AnaTrackB& candidate, AnaTrackB* vetoTrack){
//********************************************************************

    double startposz = candidate.PositionStart[2];
    double endposz = candidate.PositionEnd[2];

    if(endposz > 1850.0 && endposz > startposz
            && anaUtils::InFiducialVolume(SubDetId::kFGD1, candidate.PositionStart)){
        return false;
    }

    double TPCVetoDistance = -150.0;
    if(vetoTrack){
        if(vetoTrack->Original){
            if(vetoTrack->Original->PositionStart[2] - candidate.PositionStart[2] < TPCVetoDistance){
                return false;
            }
        }
    }
    return true;
}


//**************************************************
void cutUtils::FindOOFVtrack(const AnaEventB& event,  AnaBoxB& box,const AnaTrackB& mainTrack, const SubDetId::SubDetEnum det){
//**************************************************
    AnaTrackB* allFgdTracks[NMAXTRACKS];
    int nFgdTracks = anaUtils::GetAllTracksUsingFGDAndNoTPC(event, allFgdTracks);
    //loop over isoFGD tracks
    for (Int_t i=0;i<nFgdTracks;i++){
        AnaTrackB* track = allFgdTracks[i];

        AnaFgdTrackB *fgdTrack = track->FGDSegments[0];

        if( !fgdTrack ) continue;

        Float_t* StartPos;

        Float_t dist1 = GetSeparationSquared(fgdTrack->PositionStart, mainTrack.PositionStart);
        Float_t dist2 = GetSeparationSquared(fgdTrack->PositionEnd,   mainTrack.PositionStart);

        Float_t distance = TMath::Min(dist1,dist2); 
        if( distance == dist1 ) StartPos = track->PositionEnd;
        else StartPos = track->PositionStart;/// check that the start position of the isoFGD is well defined
        if(SubDetId::GetDetectorUsed(fgdTrack->Detector,det) && !cutUtils::FiducialCut(StartPos,det) ){
            box.OOFV = true;
            box.OOFVtrack = track;
            break;
        }
    }
}
//**************************************************
void cutUtils::FindFGD1OOFVtrack(const AnaEventB& event, AnaBoxB& box){
//**************************************************
    /// Get the OOFV FGD track
    if (box.HMNtrack) {
        FindOOFVtrack( event,box,(*box.HMNtrack),SubDetId::kFGD1 );
    }
}

//**************************************************
void cutUtils::FindFGD2OOFVtrack(const AnaEventB& event, AnaBoxB& box){
//**************************************************
    /// Get the OOFV FGD track
    if (box.HMNtrack) {
        FindOOFVtrack( event,box,(*box.HMNtrack),SubDetId::kFGD2 );
    }
}

//**************************************************
void cutUtils::FindFGD1OOFVtrack_antinu(const AnaEventB& event, AnaBoxB& box){
//**************************************************

    /// Get the OOFV FGD track
    if (box.HMPtrack) {
        FindOOFVtrack( event,box,(*box.HMPtrack),SubDetId::kFGD1 );
    }

}

//**************************************************
void cutUtils::FindFGD2OOFVtrack_antinu(const AnaEventB& event, AnaBoxB& box){
//**************************************************

    /// Get the OOFV FGD track
    if (box.HMPtrack) {
        FindOOFVtrack( event,box,(*box.HMPtrack),SubDetId::kFGD2 );
    }

}

//********************************************************************
bool cutUtils::FindLeadingTracks(const AnaEventB& event, AnaBoxB& box, bool qualityCut, SubDetId::SubDetEnum det, bool fiducialCut){
//********************************************************************


    // Get the highest momentum track (HM),
    // the highest momentum negative track (HMN) and the 
    // highest momentum positive track (HMP)
    // from all tracks passing the quality and fiducial cut. 

    // this function cannot be in anaUtils because it needs the cuts

    AnaTrackB* HM  = NULL;
    AnaTrackB* SHM  = NULL;
    AnaTrackB* HMN = NULL;
    AnaTrackB* SHMN = NULL;
    AnaTrackB* HMP = NULL;
    AnaTrackB* SHMP = NULL;

    AnaTrackB* negativeTracks[NMAXTRACKS];
    AnaTrackB* positiveTracks[NMAXTRACKS];
    int nNegativeTracks = 0;
    int nPositiveTracks = 0;

    AnaTrackB* selTracks[NMAXTRACKS];
    int nTPC = anaUtils::GetAllTracksUsingTPC(event, selTracks);

    //loop over tpc tracks
    for (Int_t i=0;i<nTPC; ++i){
        AnaTrackB* track = selTracks[i];

        if (qualityCut){
            //apply the quality cut
            if (!TrackQualityCut(*track)){
                continue;
            }
        }

        // Apply the fiducial cut
        if (fiducialCut == true && ! FiducialCut(*track, det)){
            continue;
        }

        if (track->Charge == -1){
            negativeTracks[nNegativeTracks] = track;
            nNegativeTracks++;
        }
        else{
            positiveTracks[nPositiveTracks] = track;
            nPositiveTracks++;
        }

        //tracks are sorted by momentum, so the first track is the most energetic one
        if (!HM) HM = track;
        else if (HM && !SHM) SHM = track;
        //check the charge of the most energetic track
        if (track->Charge == -1 && !HMN) {

            HMN = track;
        }

        else if (track->Charge == -1 && HMN && !SHMN) SHMN = track;

        if (track->Charge == +1 && !HMP) HMP = track;
        else if (track->Charge == +1 && HMP && !SHMP) SHMP = track;


    }

    // Fill the box
    box.HMtrack   = HM;
    box.HMNtrack  = HMN;
    box.SHMtrack  = SHM;
    box.SHMNtrack = SHMN;
    box.HMPtrack  = HMP;
    box.SHMPtrack = SHMP;
    if(event.GetIsMC() && event.SystematicsBox){
        event.SystematicsBox->HMNtrack = HMN;
        event.SystematicsBox->HMPtrack = HMP;
        event.SystematicsBox->SHMNtrack = SHMN;
    }

    for(int k = 0; k < nNegativeTracks; ++k){
        box.NegativeTPCtracks[k] = negativeTracks[k];
    }
    box.nNegativeTPCtracks = nNegativeTracks;
    for(int k = 0; k < nPositiveTracks; ++k){
        box.PositiveTPCtracks[k] = positiveTracks[k];
    }
    box.nPositiveTPCtracks = nPositiveTracks;

    return (HM);  
}

//********************************************************************
void cutUtils::FindFGDMichelElectrons(const AnaEventB& event, AnaBoxB& box){
//********************************************************************

    AnaFgdTimeBinB* arr1[NMAXFGDTIMEBINS];
    AnaFgdTimeBinB* arr2[NMAXFGDTIMEBINS];

    box.nFGD1MichelElectrons = anaUtils::GetFGDMichelElectrons(event, SubDetId::kFGD1, arr1);
    box.nFGD2MichelElectrons = anaUtils::GetFGDMichelElectrons(event, SubDetId::kFGD2, arr2);  

    for(int i = 0; i < box.nFGD1MichelElectrons; ++i){
        box.FGD1MichelElectrons[i] = arr1[i];
    }
    for(int i = 0; i < box.nFGD2MichelElectrons; ++i){
        box.FGD2MichelElectrons[i] = arr2[i];
    }

}

//**************************************************
bool cutUtils::CommonVertexCut(const AnaTrackB& track1, const AnaTrackB& track2, int cutx, int cuty, int cutz){
//**************************************************

    if (fabs(track1.PositionStart[0]-track2.PositionStart[0]) > cutx ) return false; 
    if (fabs(track1.PositionStart[1]-track2.PositionStart[1]) > cuty ) return false; 
    if (fabs(track1.PositionStart[2]-track2.PositionStart[2]) > cutz ) return false;   

    return true;
}

//**************************************************
bool cutUtils::NumuBkgMuonPIDCut(const AnaTrackB& track){
//**************************************************

    Float_t cut1 = 0.7;
    Float_t cut2 = 0.1;
    Float_t cut3 = 0.8;
    Float_t pcut = 500;

    Float_t PIDLikelihood[4];
    anaUtils::GetPIDLikelihood(track, PIDLikelihood);

    if (((PIDLikelihood[0]+PIDLikelihood[3])/(1-PIDLikelihood[2]) > cut1
                || track.Momentum > pcut ) && (PIDLikelihood[0]>cut2 &&
                    PIDLikelihood[0]<cut3)){
        return true;
    }

    return false;
}
